import numpy as np
import matplotlib.pyplot as plt
import MDSplus
import scipy.io
from scipy import signal
import pywt


def download_data(shot_num, node_name, time_context='-1:3:1e-5', server_path='192.168.20.11::/media/ennfusion/trees/exl50u'):
    """下载MDSplus数据"""
    # tree = MDSplus.Tree('exl50u', shot_num, path=server_path)
    # tree.setTimeContext(*time_context.split(':'))
    # val = tree.getNode(f'\\{node_name.upper()}').data()
    # time = tree.getNode(f'\\{node_name.upper()}').dim_of().data()

    file_name = 'sjj1.mat'
    data = scipy.io.loadmat(file_name)
    # indexs = random.sample(range(1, 51), 4)  # 随机抽取4组数据

    ii = 22
    y_original = data[f'fai{ii}'].flatten()
    y_true = data[f'ne{ii}'].flatten()

    val = y_original
    time = range(len(y_original))
    return val, time


def itoh_unwrap(phase):
    """
    Itoh算法 - 一维相位解缠绕的经典算法
    通过检测相邻点的相位差来识别跳变
    """
    unwrapped = np.zeros_like(phase)
    unwrapped[0] = phase[0]
    
    # 计算相邻点的相位差
    diff = np.diff(phase)
    
    # 检测跳变点 (通常设置为π)
    jumps = np.zeros_like(diff)
    jumps[diff < -np.pi] = 2 * np.pi
    jumps[diff > np.pi] = -2 * np.pi
    
    # 累积跳变修正
    correction = np.cumsum(jumps)
    unwrapped[1:] = phase[1:] + correction
    
    return unwrapped


def kalman_unwrap(phase, process_noise=0.01, measurement_noise=0.1):
    """
    卡尔曼滤波器解缠绕
    使用卡尔曼滤波器预测相位变化，更适合噪声较大的情况
    """
    unwrapped = np.zeros_like(phase)
    unwrapped[0] = phase[0]
    
    # 卡尔曼滤波器状态
    x_est = phase[0]  # 状态估计
    p_est = 1.0       # 估计误差协方差
    
    for i in range(1, len(phase)):
        # 预测步骤
        x_pred = x_est
        p_pred = p_est + process_noise
        
        # 计算相位差
        diff = phase[i] - (x_pred % (2 * np.pi))
        
        # 调整差值到[-π,π]范围
        if diff > np.pi:
            diff -= 2 * np.pi
        elif diff < -np.pi:
            diff += 2 * np.pi
        
        # 更新步骤
        k = p_pred / (p_pred + measurement_noise)  # 卡尔曼增益
        x_est = x_pred + k * diff
        p_est = (1 - k) * p_pred
        
        unwrapped[i] = x_est
    
    return unwrapped


def wavelet_denoise_unwrap(phase, wavelet='db4', level=3):
    """
    小波去噪+解缠绕
    先使用小波变换去除噪声，再进行相位解缠绕
    """
    # 小波去噪
    coeffs = pywt.wavedec(phase, wavelet, level=level)
    
    # 阈值处理
    for i in range(1, len(coeffs)):
        # 使用软阈值
        thresh = np.std(coeffs[i]) * np.sqrt(2 * np.log(len(coeffs[i])))
        coeffs[i] = pywt.threshold(coeffs[i], thresh, 'soft')
    
    # 重构信号
    denoised = pywt.waverec(coeffs, wavelet)
    denoised = denoised[:len(phase)]  # 确保长度一致
    
    # 使用Itoh算法解缠绕
    return itoh_unwrap(denoised)


def quality_guided_unwrap(phase, quality_map=None):
    """
    质量引导的相位解缠绕
    根据信号质量优先处理高质量区域
    """
    if quality_map is None:
        # 如果没有提供质量图，使用相位梯度的倒数作为质量指标
        dx = np.abs(np.diff(phase, append=phase[-1]))
        quality_map = 1.0 / (dx + 0.01)  # 避免除零
    
    unwrapped = np.copy(phase)
    height = len(phase)
    
    # 创建已处理标记
    processed = np.zeros(height, dtype=bool)
    
    # 从最高质量点开始
    start_idx = np.argmax(quality_map)
    processed[start_idx] = True
    
    # 创建待处理队列
    queue = [start_idx]
    
    while queue:
        current = queue.pop(0)
        
        # 检查相邻点
        neighbors = []
        if current > 0:
            neighbors.append(current - 1)
        if current < height - 1:
            neighbors.append(current + 1)
        
        for neighbor in neighbors:
            if not processed[neighbor]:
                # 计算相位差
                diff = unwrapped[neighbor] - unwrapped[current]
                
                # 调整差值
                if diff > np.pi:
                    unwrapped[neighbor] -= 2 * np.pi * np.ceil((diff - np.pi) / (2 * np.pi))
                elif diff < -np.pi:
                    unwrapped[neighbor] += 2 * np.pi * np.ceil((-diff - np.pi) / (2 * np.pi))
                
                processed[neighbor] = True
                queue.append(neighbor)
                
                # 按质量排序队列
                queue.sort(key=lambda x: -quality_map[x])
    
    return unwrapped


def process_hcn_signal(shot_num, node_name='temp01', method='itoh'):
    """
    处理HCN干涉仪信号
    
    参数:
    - shot_num: 炮号
    - node_name: 节点名称
    - method: 解缠绕方法 ('itoh', 'kalman', 'wavelet', 'quality')
    
    返回:
    - time: 时间数组
    - density: 处理后的电子密度
    """
    # 1. 下载数据
    val, time = download_data(shot_num, node_name)
    
    # 2. 预处理 - 将信号转换为相位
    # 假设信号已经是相位形式，或者需要进行适当转换
    phase_signal = val
    
    # 3. 选择相位解缠绕方法
    if method == 'itoh':
        unwrapped_phase = itoh_unwrap(phase_signal)
    elif method == 'kalman':
        unwrapped_phase = kalman_unwrap(phase_signal)
    elif method == 'wavelet':
        unwrapped_phase = wavelet_denoise_unwrap(phase_signal)
    elif method == 'quality':
        unwrapped_phase = quality_guided_unwrap(phase_signal)
    else:
        raise ValueError(f"未知的解缠绕方法: {method}")
    
    # 4. 转换为电子密度
    # 使用标准公式将相位转换为线平均电子密度
    # 假设常数: 波长=195μm, 干涉仪路径长度=35.5cm
    wavelength = 195e-6  # 米
    path_length = 0.355  # 米
    
    # 计算基线
    baseline = np.mean(unwrapped_phase[:1000])
    
    # 计算电子密度 (m^-3)
    # ne = Δφ * λ / (2π * re * L)
    # 其中re是经典电子半径 (2.82×10^-15 m)
    re = 2.82e-15  # 经典电子半径
    density = (unwrapped_phase - baseline) * wavelength / (2 * np.pi * re * path_length)
    
    return time, density


def compare_methods(shot_num, node_name='temp01'):
    """比较不同解缠绕方法的效果"""
    methods = ['itoh', 'kalman', 'wavelet', 'quality']
    
    plt.figure(figsize=(12, 10))
    
    for i, method in enumerate(methods):
        time, density = process_hcn_signal(shot_num, node_name, method)
        
        plt.subplot(len(methods), 1, i+1)
        plt.plot(time, density)
        plt.title(f'{method.capitalize()} 方法')
        plt.ylabel('电子密度 (m^-3)')
        
        if i == len(methods) - 1:
            plt.xlabel('时间 (s)')
    
    plt.tight_layout()
    plt.savefig(f'e:\\ENN\\2025工作\\诊断-hcn拼接\\hcn_comparison_{shot_num}.png')
    plt.show()


if __name__ == "__main__":
    # 示例: 处理炮号4417的数据
    shot = 4417
    
    # 比较不同方法
    compare_methods(shot)
    
    # 使用最佳方法处理并绘图
    time, density = process_hcn_signal(shot, method='wavelet')
    
    plt.figure(figsize=(10, 6))
    plt.plot(time, density)
    plt.xlabel('时间 (s)')
    plt.ylabel('电子密度 (m^-3)')
    plt.title('HCN干涉仪电子密度 - 小波解缠绕方法')
    plt.grid(True)
    plt.savefig(f'e:\\ENN\\2025工作\\诊断-hcn拼接\\hcn_density_{shot}.png')
    plt.show()
    
    print('处理完成!')